﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Text;

namespace SecurityModule
{
    class ReplaceFilter : Stream
    {
        //Temporary buffer to accumulate page writes
        private readonly MemoryStream _ms;

        //Handle to original output stream pipeline
        private readonly Stream _outputStream;

        //Encoding of the response
        private readonly Encoding _encoding;

        //The callback function
        private readonly Func<string, string> _callback;
        private readonly Func<Stream, Encoding, Stream> _streamingCallback;

        /// <summary>
        /// String-based callback filter.
        /// </summary>
        /// <param name="stream">The stream to filter.</param>
        /// <param name="encoding">The <see cref="Encoding"/> to be used to decode the stream to a string.</param>
        /// <param name="callback">A function that takes <paramref name="stream"/> contents as a string and returns a replacement for the whole thing.</param>
        public ReplaceFilter(Stream stream, Encoding encoding, Func<string, string> callback)
        {
            _ms = new MemoryStream();
            _outputStream = stream;
            _encoding = encoding;
            _callback = callback;
        }
        /// <summary>
        /// String-based callback filter using a dictionary for replacements.
        /// </summary>
        /// <param name="stream">The stream to filter.</param>
        /// <param name="encoding">The <see cref="Encoding"/> to be used to decode the stream to a string.</param>
        /// <param name="replacements">A dictionary defining text to replace where the occarnces of keys are replaced with corresponding values.</param>
        public ReplaceFilter(Stream stream, Encoding encoding, Dictionary<string, string> replacements)
        {
            _ms = new MemoryStream();
            _outputStream = stream;
            _encoding = encoding;
            _callback = content =>
                {
                    var sb = new StringBuilder(content);
                    foreach (var replacement in replacements)
                        sb.Replace(replacement.Key, replacement.Value);
                    return sb.ToString();
                };
        }

        /// <summary>
        /// Streaming callback filter.
        /// </summary>
        /// <param name="stream">The stream to filter.</param>
        /// <param name="encoding">The <see cref="Encoding"/> to be used to decode the stream to a string.</param>
        /// <param name="callback">A function that takes <paramref name="stream"/> contents and <paramref name="encoding"/> and returns a replacement for the whole stream.</param>
        public ReplaceFilter(Stream stream, Encoding encoding, Func<Stream, Encoding, Stream> callback)
        {
            _ms = new MemoryStream();
            _outputStream = stream;
            _encoding = encoding;
            _streamingCallback = callback;
        }

        public override void Flush()
        {
            //Decide if streaming or string based filtering

            if (_streamingCallback != null)
            {
                _ms.Seek(0, SeekOrigin.Begin);
                var stream = _streamingCallback(_ms, _encoding);
                stream.Seek(0, SeekOrigin.Begin);
                stream.CopyTo(_outputStream);
                _outputStream.Flush();
            }
            else if (_callback != null)
            {
                //perform replacement
                var content = _encoding.GetString(_ms.GetBuffer());

                //Perform the content replacement routine
                content = _callback(content);

                var bytes = _encoding.GetBytes(content);
                _outputStream.Write(bytes, 0, bytes.Length);
                _outputStream.Flush();
            }



            //This is a passthrough filter
            //ms.Seek(0, SeekOrigin.Begin);
            //ms.CopyTo(_outputStream);
            //_outputStream.Flush();
        }

        public override long Seek(long offset, SeekOrigin origin)
        {
            return _ms.Seek(offset, origin);
        }

        public override void SetLength(long value)
        {
            _ms.SetLength(value);
        }

        public override int Read(byte[] buffer, int offset, int count)
        {
            return _ms.Read(buffer, offset, count);
        }

        public override void Write(byte[] buffer, int offset, int count)
        {

            _ms.Write(buffer, offset, count);
            //_outputStream.Write(buffer, offset, count);

        }

        public override bool CanRead
        {
            get { return true; }
        }

        public override bool CanSeek
        {
            get { return true; }
        }

        public override bool CanWrite
        {
            get { return true; }
        }

        public override long Length
        {
            get { return _ms.Length; }
        }

        public override long Position
        {
            get { return _ms.Position; }
            set { _ms.Position = value; }
        }
    }
}
